import torch
import torch.nn as nn

from Modules.SoftMax2dArea import SoftMax2dArea
from Modules.MapConv import MapConv


class SoftPool(nn.Module):
    def __init__(self, in_channel, hidden_channel=64, adapt_pool_size=7):
        super().__init__()
        self.map_conv = MapConv(in_channel, hidden_channel, adapt_pool_size)
        self.pool = nn.AdaptiveAvgPool2d(1)

    def forward(self, x):
        b,c,h,w = x.size()
        map = self.map_conv(x)
        x = map * x
        x = self.pool(x) * h * w / 2
        return x


class SoftPooling1D(torch.nn.Module):
    def __init__(self, kernel_size, strides=None, padding=0, ceil_mode=False, count_include_pad=True):
        super(SoftPooling1D, self).__init__()
        self.avgpool = torch.nn.AvgPool1d(kernel_size, strides, padding, ceil_mode, count_include_pad)

    def forward(self, x):
        x_exp = torch.exp(x)
        x_exp_pool = self.avgpool(x_exp)
        x = self.avgpool(x_exp * x)
        return x / x_exp_pool


class SoftPooling2D(torch.nn.Module):
    def __init__(self, kernel_size, strides=None, padding=0, ceil_mode=False, count_include_pad=True,
                 divisor_override=None):
        super(SoftPooling2D, self).__init__()
        self.avgpool = torch.nn.AvgPool2d(kernel_size, strides, padding, ceil_mode, count_include_pad, divisor_override)

    def forward(self, x):
        x_exp = torch.exp(x)
        x_exp_pool = self.avgpool(x_exp)
        x = self.avgpool(x_exp * x)
        return x / x_exp_pool


class SoftPooling3D(torch.nn.Module):
    def __init__(self, kernel_size, strides=None, padding=0, ceil_mode=False, count_include_pad=True,
                 divisor_override=None):
        super(SoftPooling3D, self).__init__()
        self.avgpool = torch.nn.AvgPool3d(kernel_size, strides, padding, ceil_mode, count_include_pad, divisor_override)

    def forward(self, x):
        x_exp = torch.exp(x)
        x_exp_pool = self.avgpool(x_exp)
        x = self.avgpool(x_exp * x)
        return x / x_exp_pool


if __name__ == '__main__':
    # Example Usage:
    x = torch.randn(2, 64,7,7)

    model = SoftPooling2D(7,1)
    y = model(x)
    print(y.size())



